## 卷积神经网络识别轴承数据

# 数据处理和加载

# 这部分主要用于从`.mat`文件中加载数据，并将数据划分为训练集和测试集。
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score,confusion_matrix,classification_report
import matplotlib.pyplot as plt
import seaborn as sns
#from tensorboardX import SummaryWriter
import copy
import time
from pylab import mpl
import torch
import torch.nn as nn
from torch.optim import Adam
import torch.utils.data as Data
from torchvision import transforms

import scipy.io as scio

# 加载、处理、划分数据
''' 导入数据 '''
 # 源域-训练集
dataFile = r'C:\Users\viki6\Desktop\研一\实验室\研一上\行星轮系故障诊断测试\data_target\1#\X_train.mat'
data = scio.loadmat(dataFile)
x_train_s = data['X_train']
x_train_s = np.transpose(x_train_s)
x_train_s = x_train_s.reshape(-1, 1, 8192) 

dataFile = r'C:\Users\viki6\Desktop\研一\实验室\研一上\行星轮系故障诊断测试\data_target\1#\train_Y.mat'
data = scio.loadmat(dataFile)
y_train_s = data['train_Y']
y_train_s = np.squeeze(y_train_s)
y_train_s = y_train_s - 1


# 源域-测试集
dataFile = r'C:\Users\viki6\Desktop\研一\实验室\研一上\行星轮系故障诊断测试\data_target\1#\X_test.mat'
data = scio.loadmat(dataFile)
x_test = data['X_test']
x_test = np.transpose(x_test)
x_test = x_test.reshape(-1, 1, 8192)       # reshape
       
dataFile = r'C:\Users\viki6\Desktop\研一\实验室\研一上\行星轮系故障诊断测试\data_target\1#\test_Y.mat'
data = scio.loadmat(dataFile)
y_test = data['test_Y']
y_test = np.squeeze(y_test)
y_test = y_test - 1 

x_train_s1 = torch.from_numpy(x_train_s).type(torch.FloatTensor)  
y_train_s1 = torch.from_numpy(y_train_s).type(torch.LongTensor) 

x_test = torch.from_numpy(x_test).type(torch.FloatTensor)
y_test = torch.from_numpy(y_test).type(torch.LongTensor) 

sourcedata = Data.TensorDataset(x_train_s1, y_train_s1)
testdata = Data.TensorDataset(x_test, y_test)

# 使用DataLoader加载数据
train_loader = Data.DataLoader(dataset=sourcedata, batch_size=128, shuffle=True, drop_last=True)
test_data = Data.DataLoader(dataset=testdata, batch_size=128, shuffle=False, drop_last=True)


from torch import nn
import warnings

# 定义网络结构
# 网络由五个卷积层和一个全连接层组成
#CNN 类继承了 PyTorch 的基础 nn.Module 类，该类是所有神经网络模块的基础类，可以自定义你自己的神经网络模型。
class CNN(nn.Module): #定义了一个名为 CNN 的类，继承了 PyTorch 的 nn.Module 类。这样 CNN 就具备了 nn.Module 类里面的所有功能和方法。
    def __init__(self, pretrained=False, in_channel=1, out_channel=4): #定义了 CNN 类的构造函数，这个函数会在你创建这个类的对象的时候被调用。
        # 函数的参数有三个：pretrained 参数用来指定是否使用预训练模型；in_channel 指定了第一层卷积层的输入通道的数量；out_channel 指定了最终的输出通道的数量。
        super(CNN, self).__init__() #这行代码调用了它父类 nn.Module 的构造方法。这对于 PyTorch 模型是必须的，如果没有这行代码，那么你的网络就不包含任何层，它也不会知道有哪些参数需要学习。
        if pretrained == True: #如果 pretrained 参数为 True，它将发出一个警告，表示预训练模型不可用。这在你希望使用预训练的模型，但预训练模型并未提供的情况下，会非常有用。
            warnings.warn("Pretrained model is not available")

        #这段代码定义了第一个卷积层（包含一个卷积层，一个批标准化层，一个ReLU函数和一个最大池化层）
        self.layer1 = nn.Sequential(
            nn.Conv1d(in_channel, 32, kernel_size = 64, stride = 1, padding = 32),
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2))
        #这段代码定义了第二个卷积层（同样包含一个卷积层，一个批标准化层，一个 ReLU 函数和一个最大池化层）
        self.layer2 = nn.Sequential(
            nn.Conv1d(32, 48, kernel_size = 16, stride = 1, padding = 8),
            nn.BatchNorm1d(48),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2))

        self.layer3 = nn.Sequential(
            nn.Conv1d(48, 64, kernel_size = 5, stride = 1, padding = 2),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2))

        self.layer4 = nn.Sequential(
            nn.Conv1d(64, 64, kernel_size = 5, stride = 1, padding = 2),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2))

        self.layer5 = nn.Sequential(
            nn.Linear(64 * 64 * 8, 512),
            nn.ReLU(inplace=True))
        self.fc = nn.Linear(512, out_channel)

    #这段代码定义了前向传播过程和创建了一个CNN模型的实例
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        #这些语句依次通过前面的定义的四层卷积神经网络
        x = x.view(x.size(0), -1)
        #这句代码用于将卷积输出的多维数据拉平为一维，因为全连接层只接收一维数据
        x = self.layer5(x)
        x = self.fc(x)
        #这两行代码分别代表一层卷积层和一个全连接层的处理过程。全连接层self.fc的目标通常是将前面的高维特征向量映射到目标类别上

        return x  #最后返回了网络模型的输出，这个输出将用于计算损失

myconvnet = CNN()

## 卷积神经网络训练与预测
# 定义网络的训练过程

# 训练模型函数
def train_model(model,traindataloader, train_rate, criterion,
                optimizer, num_epochs=25):
    """
    model:网络模型；traindataloader:训练数据集，会切分为训练集和验证集
    train_rate:训练集batchsize百分比;criterion：损失函数；optimizer：优化方法；
    num_epochs:训练的轮数
    """
    ## 计算训练使用的batch数量
    batch_num = len(traindataloader)
    train_batch_num = round(batch_num*train_rate)
    ## 复制模型的参数
    best_model_wts = copy.deepcopy(model.state_dict())  # 先深拷贝一份当前模型的参数，后面迭代过程中若遇到更优模型则替换。
    best_acc = 0.0
    train_loss_all = []
    train_acc_all = []
    val_loss_all = []
    val_acc_all = []
    since = time.time()
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        print('-'*10)
        ## 每个epoch有两个训练阶段
        train_loss = 0.0
        train_corrects = 0
        train_num = 0
        val_loss = 0.0
        val_corrects = 0
        val_num = 0

        #这段代码是一个神经网络模型的训练和验证过程。

        for step,(b_x,b_y) in enumerate(traindataloader):  #这一行代码遍历了训练数据 traindataloader，每次循环都会得到一个批次的数据 (b_x 为输入特征，b_y 为正确标签) 和这个批次的索引 step。
            #模型训练部分：
            if step < train_batch_num:
                model.train()   # 设置模型为训练模式
                # b_x = b_x.long()
                output = model(b_x) # 将输入数据传入模型中，进行前向传播获取输出
                pre_lab = torch.argmax(output, 1) # 获取预测结果，torch.argmax函数可以返回每行中最大元素的索引，也就是预测的类别
                loss = criterion(output, b_y.long()) # 计算预测结果和实际标签之间的损失/误差，通常是使用交叉熵作为损失函数
                optimizer.zero_grad() # 优化器梯度清零，因为pytorch会累积梯度
                loss.backward() # 计算loss的反向传播
                optimizer.step() # 根据反向传播求得的梯度进行参数更新
                train_loss += loss.item() * b_x.size(0)  # 计算总训练损失，.item()是获取loss的具体值
                train_corrects += torch.sum(pre_lab == b_y.data) # 计算预测正确的样本数，预测结果和实际结果进行对比，结果为True的即为预测正确的样本
                train_num += b_x.size(0) # 计算总的样本数
            #验证训练模型的效果：
            else:
                model.eval()  ## 设置模型为训练模式评估模式
                output = model(b_x)
                pre_lab = torch.argmax(output, 1)
                loss = criterion(output, b_y.long())
                val_loss += loss.item() * b_x.size(0)  # 计算总验证损失
                val_corrects += torch.sum(pre_lab == b_y.data)   # 计算预测正确的样本数
                val_num += b_x.size(0)  # 计算总的样本数
        ## 计算一个epoch在训练集和验证集上的的损失和精度
        train_loss_all.append(train_loss / train_num)
        train_acc_all.append(train_corrects.double().item() / train_num)
        val_loss_all.append(val_loss / val_num)
        val_acc_all.append(val_corrects.double().item() / val_num)
        print('{} Train Loss: {:.4f}  Train Acc: {:.4f}'.format(
            epoch, train_loss_all[-1], train_acc_all[-1]))
        print('{} Val Loss: {:.4f}  val Acc: {:.4f}'.format(
            epoch, val_loss_all[-1], val_acc_all[-1]))
        # 拷贝模型最高精度下的参数
        if val_acc_all[-1] > best_acc:
            best_acc = val_acc_all[-1]
            best_model_wts = copy.deepcopy(model.state_dict())
        time_use = time.time() - since
        print("Train and val complete in {:.0f}m {:.0f}s".format(
            time_use // 60, time_use % 60))
    # 使用最好模型的参数
    model.load_state_dict(best_model_wts)
    train_process = pd.DataFrame(
        data={"epoch": range(num_epochs),
                "train_loss_all": train_loss_all,
                "val_loss_all": val_loss_all,
                "train_acc_all": train_acc_all,
                "val_acc_all": val_acc_all})
    return model, train_process

## 对模型进行训练
optimizer = torch.optim.Adam(myconvnet.parameters(), lr=0.0003)
# 定义损失函数，使用交叉熵损失函数
criterion = nn.CrossEntropyLoss()   # 损失函数
# 切换模型为训练模式
myconvnet,train_process = train_model(
    myconvnet,train_loader, 0.8,
    criterion, optimizer, num_epochs=12)

## 可视化模型训练过程中
plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plt.plot(train_process.epoch,train_process.train_loss_all,
         "ro-",label = "Train loss")
plt.plot(train_process.epoch,train_process.val_loss_all,
         "bs-",label = "Val loss")
plt.legend()
plt.xlabel("epoch")
plt.ylabel("Loss")
plt.subplot(1,2,2)
plt.plot(train_process.epoch,train_process.train_acc_all,
         "ro-",label = "Train acc")
plt.plot(train_process.epoch,train_process.val_acc_all,
         "bs-",label = "Val acc")
plt.xlabel("epoch")
plt.ylabel("acc")
plt.legend()
plt.show()

## 对测试集进行预测，并可视化预测效果
torch.save(myconvnet,'net1.pth')
torch.save(myconvnet.state_dict(),'net1_params.pth')

#保存训练好的整个网络：torch.save(net1,'net1.pth')
#只保存网络的参数：torch.save(net1.state_dict(),'net1_params.pth')

#调用整个训练好的网络：net2 = torch.load('net1.pth')
#只调用网络的参数：net3.load_state_dict(torch.load('net1_params.pth'))

myconvnet.eval()
# 得到模型的输出结果
output = myconvnet(x_test)
# 得到模型预测的标签
pre_lab = torch.argmax(output,1)
# 计算模型预测的准确率
acc = accuracy_score(y_test,pre_lab)
# 输出准确率
print("在测试集上的预测精度为:",acc)

class_label = ['太阳轮点蚀', '太阳轮裂纹', '行星轮点蚀', '行星轮裂纹']

# print(y_test.shape)
# print(pre_lab.shape)
## 计算混淆矩阵并可视化
plt.figure(figsize=(10,8))
# 对预测结果进行展示（这里展示了混淆矩阵）
conf_mat = confusion_matrix(y_test,pre_lab)
df_cm = pd.DataFrame(conf_mat, index=class_label,
                     columns=class_label)
sns.set(font_scale=1.5)
heatmap = sns.heatmap(df_cm, annot=True, fmt="d",cmap="YlGnBu")
heatmap.yaxis.set_ticklabels(heatmap.yaxis.get_ticklabels(), rotation=0, ha='right')
heatmap.xaxis.set_ticklabels(heatmap.xaxis.get_ticklabels(), rotation=20, ha='right')
font2 = {'family' : 'SimHei',    'weight' : 'normal',    'size' : 24,}
# 绘制混淆矩阵
plt.ylabel('真实标签', font2)
plt.xlabel('预测标签', font2)
plt.show()

# plt.rcParams两行是用于解决标签不能显示汉字的问题
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus'] = False

mpl.rcParams["font.sans-serif"] = ["SimHei"]
mpl.rcParams["axes.unicode_minus"] = False


torch.set_printoptions(profile=True)
abc = open('parameter1.txt','w')
para = list(myconvnet.parameters())
#print(para)
#print(para,file=abc)
abc.close()